package com.aiwiown.snackmq.broker.handler;

import com.aiwiown.snackmq.broker.auth.AuthorizationService;
import com.aiwiown.snackmq.cluster.raft.SnackMQStateMachine;
import com.aiwiown.snackmq.cluster.raft.entity.SnackMQRaftLogEntry;
import com.aiwiown.snackmq.cluster.raft.manager.SnackMQRaftNodeManager;
import com.aiwiown.snackmq.common.auth.Action;
import com.aiwiown.snackmq.common.exception.AuthorizationException;
import com.aiwiown.snackmq.common.handler.MessageHandler;
import com.aiwiown.snackmq.common.handler.RequestContext;
import com.aiwiown.snackmq.common.message.Message;
import com.aiwiown.snackmq.common.message.MessageStatus;
import com.aiwiown.snackmq.common.message.MessageType;
import com.aiwiown.snackmq.common.protocol.OperationType;
import com.aiwiown.snackmq.common.protocol.TopicPartition;
import com.aiwiown.snackmq.storage.transaction.TransactionLog;
import com.aiwiown.snackmq.storage.transaction.TransactionRecord;
import com.aiwiown.snackmq.storage.transaction.TransactionState;
import lombok.extern.slf4j.Slf4j;

import java.io.IOException;
import java.util.Collections;

/**
 * 处理事务消息的第二阶段：提交事务。
 */
@Slf4j
// 【修改】移除 @RequiredArgsConstructor 以便使用自定义构造函数
public class CommitMessageHandler implements MessageHandler {

    private final TransactionLog transactionLog;
    private final SnackMQRaftNodeManager dataRaftNodeManager;
    private final AuthorizationService authorizationService;
    // --- 【新增】添加 StateMachine 依赖 ---
    private final SnackMQStateMachine stateMachine;

    /**
     * 【新增】构造函数，注入 SnackMQStateMachine。
     */
    public CommitMessageHandler(TransactionLog transactionLog,
                                SnackMQRaftNodeManager dataRaftNodeManager,
                                AuthorizationService authorizationService,
                                SnackMQStateMachine stateMachine) {
        this.transactionLog = transactionLog;
        this.dataRaftNodeManager = dataRaftNodeManager;
        this.authorizationService = authorizationService;
        this.stateMachine = stateMachine;
    }


    @Override
    public MessageType getMessageType() {
        return MessageType.COMMIT_TRANSACTIONAL_MESSAGE;
    }

    @Override
    public void handle(RequestContext ctx, Message request) {
        String transactionId = request.getProperty("transactionId");
        if (transactionId == null || transactionId.isEmpty()) {
            ctx.writeResponse(createErrorResponse(request, "TransactionId is missing."));
            return;
        }

        Message halfMessage = transactionLog.getPendingMessage(transactionId);
        if (halfMessage == null) {
            log.warn("Cannot commit transaction '{}': it may have already been resolved or is invalid.", transactionId);
            ctx.writeResponse(createErrorResponse(request, "Transaction not found or already resolved."));
            return;
        }

        try {
            // 授权检查：提交事务等同于对原始主题进行生产操作
            // 【修复】使用新的 Resource 对象进行授权，明确资源类型为 TOPIC
            authorizationService.authorize(request.getProperty("authToken"), Action.PRODUCE, new AuthorizationService.Resource("TOPIC", halfMessage.getTopic()));

            // --- 【核心修复】: 在提议到 Raft 之前，为消息分配一个 offset ---
            String topic = halfMessage.getTopic();
            // 从半消息中获取分区信息，如果不存在则默认为0
            int partitionId = Integer.parseInt(halfMessage.getProperty("partition", "0"));
            TopicPartition topicPartition = new TopicPartition(topic, partitionId);

            // 从状态机持有的存储服务中获取下一个可用的 offset
            long nextOffset = stateMachine.getStorageService().getNextOffset(topic, partitionId);

            // 将获取到的 offset 添加到消息属性中，创建一个新的、完整的待提交消息
            Message messageToCommit = halfMessage.toBuilder()
                    .property("offset", String.valueOf(nextOffset))
                    .build();

            // 使用带有 offset 的消息创建 Raft 日志条目
            Message finalBatchForRaft = Message.builder()
                    .type(MessageType.PRODUCE_REQUEST)
                    .messages(Collections.singletonList(messageToCommit))
                    .build();
            SnackMQRaftLogEntry logEntry = new SnackMQRaftLogEntry(OperationType.PRODUCE, finalBatchForRaft);
            // --- 【修复结束】 ---

            dataRaftNodeManager.apply(logEntry, status -> {
                if (status.isOk()) {
                    // 主日志写入成功后，更新事务日志，将事务标记为 COMMITTED
                    try {
                        String producerId = halfMessage.getProperty("producerId");
                        TransactionRecord commitRecord = new TransactionRecord(
                                transactionId,
                                TransactionState.COMMITTED,
                                System.currentTimeMillis(),
                                null, // 提交记录不需要消息体
                                producerId // 传递 producerId
                        );
                        transactionLog.append(commitRecord);
                        ctx.writeResponse(createSuccessResponse(request));
                        log.info("Transaction {} committed successfully.", transactionId);
                    } catch (IOException e) {
                        // 这是一个严重错误。消息已写入主日志，但事务状态未更新。
                        // 这可能导致在状态回查时消息被重复投递。
                        log.error("CRITICAL: Failed to append COMMIT record for transaction {}. This may lead to message duplication.", transactionId, e);
                        ctx.writeResponse(createErrorResponse(request, "Failed to finalize transaction state. Please check broker logs."));
                    }
                } else {
                    log.error("Failed to apply committed message for transaction {} to Raft log. Status: {}", transactionId, status);
                    ctx.writeResponse(createErrorResponse(request, "Failed to commit message to main log: " + status.getErrorMsg()));
                }
            });

        } catch (AuthorizationException e) {
            log.warn("Authorization failed for commit transaction request: {}", e.getMessage());
            ctx.writeResponse(createErrorResponse(request, "Permission denied: " + e.getMessage()));
        } catch (Exception e) {
            log.error("Error processing commit transaction request for transactionId: {}", transactionId, e);
            ctx.writeResponse(createErrorResponse(request, "Internal server error."));
        }
    }

    private Message createSuccessResponse(Message request) {
        return Message.builder().messageId(request.getMessageId()).type(MessageType.RESPONSE).status(MessageStatus.SUCCESS).build();
    }

    private Message createErrorResponse(Message request, String errorMessage) {
        return Message.builder().messageId(request.getMessageId()).type(MessageType.RESPONSE).status(MessageStatus.FAILED).errorMessage(errorMessage).build();
    }
}